package com.hoe.tsb.fraft.net;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.hoe.tsb.Configuration;
import com.hoe.tsb.fraft.LRUList;
import com.hoe.tsb.fraft.Node;
import com.hoe.tsb.fraft.enums.MessageType;
import com.hoe.tsb.fraft.enums.RoleType;
import com.hoe.tsb.fraft.event.msg.FRaftMessage;
import com.hoe.tsb.fraft.exception.FRaftException;
import com.hoe.tsb.fraft.serialize.SerializerHelper;
import com.xiaoleilu.hutool.util.HexUtil;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.*;

/**
 * <p>
 * 基于UDP的数据交换器
 * </p>
 *
 * 已简单实现ACK以及重传机制
 *
 * @author hoe
 * @since 2022/11/21
 */
@Slf4j
public class UdpExchanger extends AbstractExchanger {

    static class Msg {
        byte[] originData;
        FRaftMessage body;
        String sign;

        /**
         * 接收数据时
         * @param node 发送端的节点信息
         * @param body 接收到的数据体
         */
        Msg(Node node, byte[] body) {
            this.originData = body;
            //反序列化body并构建出FRaftMessage
            FRaftMessage fRaftMessage = SerializerHelper.desWithSign(body);
            fRaftMessage.setNodes(Arrays.asList(node));
            this.body = fRaftMessage;
            //计算sign
            this.sign = genSign(originData);
        }

        /**
         * 发送数据时
         * @param body 待发送的FRaft消息
         */
        Msg(FRaftMessage body) {
            this.body = body;
            //序列化body
            this.originData = SerializerHelper.serWithSign(body);
            //计算sign
            this.sign = genSign(originData);
        }

        private String genSign(byte[] body) {
            byte[] signBytes = new byte[16];
            System.arraycopy(body, 1, signBytes, 0, 16);
            return HexUtil.encodeHexStr(signBytes);
        }
    }

    /**
     * FIFO 发送队列
     */
    private final LinkedBlockingQueue<Msg> msgSendQueue = new LinkedBlockingQueue<>();
    /**
     * FIFO 接收队列
     */
    private final LinkedBlockingQueue<Msg> msgRcvQueue = new LinkedBlockingQueue<>();
    /**
     * 等待ACK的Map<br/>
     * key：消息摘要<br/>
     * value：消息体<br/>
     * TODO 目前没有对该Map的长度进行限制，如果网络拥塞可能有OOM的风险
     */
    private final ConcurrentHashMap<String, Msg> waitAckMap = new ConcurrentHashMap<>();
    /**
     * 校验重复消息的数量是固定的（可配置），超过长度就采取LRU进行删除
     */
    private final List<String> repeatSign = new LRUList<>(1024);
    /**
     * 创建一个4核心线程的线程池，1个用来发消息，1个用来接收消息，1个用来取消息，1个用来重传，其他的用来响应消息<br/>
     * 瓶颈在msgQueue中，这里16个线程足以应对<br/>
     * 以防万一，这里的Policy还是采用CallRunner，不然还要处理消息的重回调逻辑
     */
    private Executor mExec = new ThreadPoolExecutor(4, 16,
            60L, TimeUnit.SECONDS, new SynchronousQueue<>(),
            new ThreadFactoryBuilder().setNameFormat("udp-task-%d").build(),
            new ThreadPoolExecutor.CallerRunsPolicy());

    /**
     * 发送接收UDP报文的Socket
     */
    private final DatagramSocket mUdpSocket;
    private final Configuration mConfig;

    public UdpExchanger(String ip, int port, String pubKey) {
        this(Configuration.builder().info(new Configuration.Info(ip, port, pubKey)).build());
    }

    public UdpExchanger(Configuration config) {
        super(config.getCurrent().getIp(), config.getCurrent().getPort());
        this.mConfig = config;

        try {
            mUdpSocket = new DatagramSocket(port);
        } catch (SocketException e) {
            throw new FRaftException("初始化Socket异常", e);
        }

        //启动发送线程
        mExec.execute(() -> {
            while(!willClose) {
                Msg msg;
                try {
                    msg = msgSendQueue.take();
                } catch (InterruptedException e) {
                    log.error("取发送消息任务被打断", e);
                    continue;
                }

                //实现UDP报文发送
                List<Node> nodes = msg.body.getNodes();
                if (nodes == null || nodes.size() == 0) {
                    throw new FRaftException("需要通信的节点不能为空");
                }
                nodes.forEach(n -> {
                    InetAddress iNetAddress = getINetAddress(n);
                    DatagramPacket packet =  new DatagramPacket(msg.originData, msg.originData.length, iNetAddress, n.getPort());
                    try {
                        mUdpSocket.send(packet);
                    } catch (IOException e) {
                        log.error("发送消息到ip：" + n.getIp() + " port：" + n.getPort() + " 出错了", e);
                    }
                });
            }
        });
        //启动接收线程
        //WARN：发送消息不能超过5MB（超过时将被直接截断）
        final byte[] rcvBuf = new byte[1024 * 1024 * 5];
        final DatagramPacket rcvPacket = new DatagramPacket(rcvBuf, rcvBuf.length);
        mExec.execute(() -> {
            while(!willClose) {
                try {
                    mUdpSocket.receive(rcvPacket);
                } catch (IOException e) {
                    log.error("监听UDP错误", e);
                    continue;
                }

                byte[] data = Arrays.copyOf(rcvPacket.getData(), rcvPacket.getLength());
                Node rcvNode = new Node(rcvPacket.getAddress().getHostAddress(), rcvPacket.getPort());
                Msg msg = new Msg(rcvNode, data);
                log.debug("收到UDP消息，长度：{}，摘要：{}", data.length, msg.sign);
                //剔除ACK方法
                if (msg.body.getType() == MessageType.ACK) {
                    String ackSign = new String(msg.body.getBody(), StandardCharsets.UTF_8);
                    log.debug("收到{}:{}的ACK[{}]，ACK本身的Sign[{}]", rcvNode.getIp(), rcvNode.getPort(), ackSign, msg.sign);
                    //从等待ACK队列移除数据
                    waitAckMap.remove(ackSign);
                    continue;
                }

                //回复ACK
                //TODO 此处没有动态配置ACK的策略，后续可优化成可动态配置消费成功再ACK、收到即ACK
                if (msg.body.getType().isNeedAck()) {
                    Msg ackMsg = new Msg(new FRaftMessage(MessageType.ACK, 0, RoleType.FOLLOWER, msg.sign.getBytes(StandardCharsets.UTF_8), Arrays.asList(rcvNode)));
                    msgSendQueue.add(ackMsg);
                    log.debug("对消息回复ACK，原消息摘要：{}，ACK摘要：{}", msg.sign, ackMsg.sign);
                }

                //如果已经响应过该消息，则不重复响应
                //此处可能会出现消费该消息消费失败了却无法补消费的问题（只能客户端重发请求）
                if (repeatSign.contains(msg.sign)) {
                    //TODO 消息如果不加入发送节点、时间戳，则很多重复消息，因为摘要都一样
                    log.debug("已经消费过{}:{}的消息[{}]", rcvNode.getIp(), rcvNode.getPort(), msg.sign);
                    continue;
                } else {
                    repeatSign.add(msg.sign);
                }

                try {
                    msgRcvQueue.put(msg);
                } catch (InterruptedException e) {
                    log.error("存放接收消息到缓存队列出错", e);
                    continue;
                }
            }
        });
        //启动回调线程
        mExec.execute(() -> {
            while(!willClose) {
                MessageReceive receiver = getReceiver();
                if (receiver != null) {
                    Msg msg;
                    try {
                        msg = msgRcvQueue.take();
                    } catch (InterruptedException e) {
                        log.error("取发送消息任务被打断", e);
                        continue;
                    }

                    //响应消息
                    log.info("收到来自{}:{}的消息，角色类型：{}，任期：{}，消息类型：{}", msg.body.getNodes().get(0).getIp(), msg.body.getNodes().get(0).getPort(),
                            msg.body.getRoleType().name(), msg.body.getTerm(), msg.body.getType().name());
                    mExec.execute(() -> receiver.onReceive(msg.body));
                } else {
                    try {
                        log.warn("没有设置UDP消息接收器，等待设置中...");
                        Thread.sleep(1000);
                    } catch (InterruptedException e) {
                        log.error("等待设置UDP消息接收器出错", e);
                    }
                }
            }
        });
        //启动重传线程
        mExec.execute(() -> {
            while (!willClose) {
                if (!waitAckMap.isEmpty()) {
                    log.debug("检测到需要重传的消息，开始执行重传...");
                    for (Msg msg : waitAckMap.values()) {
                        log.debug("待重传消息类型：{}，摘要：{}", msg.body.getType().name(), msg.sign);
                    }
                    msgSendQueue.addAll(waitAckMap.values());
                }
                synchronized (waitAckMap) {
                    try {
                        waitAckMap.wait(mConfig.getAckTimeout());
                    } catch (InterruptedException e) {
                        log.error("重传线程出错", e);
                    }
                }
            }
        });
    }

    @Override
    public String getName() {
        return "udp";
    }

    @Override
    public void send(FRaftMessage raftMessage) {
        //记录发送日志
        logSend(raftMessage);
        mExec.execute(() -> {
            try {
                Msg msg = new Msg(raftMessage);
                //无需调用offer实现定长阻塞方法
                //因为这里直到阻塞8个线程后会在线程池的execute方法进行溢出处理
                msgSendQueue.put(msg);
                //需要回复就加入到等待ACK队列
                //此队列会定时重传
                if (raftMessage.getType().isNeedAck()) {
                    waitAckMap.put(msg.sign, msg);
                }
            } catch (InterruptedException e) {
                log.error("往队列中存放消息错误", e);
            }
        });
    }

    @Override
    protected void onClose() {
        if (mUdpSocket != null) {
            mUdpSocket.close();
        }
    }

    private void logSend(FRaftMessage msg) {
        try {
            List<Node> nodes = msg.getNodes();
            for (Node n : nodes) {
                log.info("向{}:{}发送UDP消息，消息类型:{}", n.getIp(), n.getPort(), msg.getType().name());
            }
        } catch (NullPointerException e) {
            throw new FRaftException("发送消息时待接收节点不能为空");
        }
    }

    private InetAddress getINetAddress(Node node) {
        try {
            return InetAddress.getByName(node.getIp());
        } catch (UnknownHostException e) {
            throw new FRaftException("IP地址转换错误：" + node.getIp(), e);
        }
    }
}
